import tkinter as tk
from PIL import Image, ImageDraw
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
import matplotlib.pyplot as plt
import os

# 定义与训练时相同的模型结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

class DrawingApp:
    def __init__(self, root, model, device):
        self.root = root
        self.root.title("手写数字识别")
        
        # 设置模型和设备
        self.model = model
        self.device = device
        
        # 创建画布
        self.canvas_size = 280  # 10倍于MNIST的28x28
        self.canvas = tk.Canvas(root, width=self.canvas_size, height=self.canvas_size, bg="black")
        self.canvas.pack(side=tk.LEFT, padx=10, pady=10)
        
        # 创建图像和绘图对象
        self.image = Image.new("L", (self.canvas_size, self.canvas_size), color=0)
        self.draw = ImageDraw.Draw(self.image)
        
        # 绑定鼠标事件
        self.canvas.bind("<B1-Motion>", self.paint)
        
        # 创建控制面板
        self.control_frame = tk.Frame(root)
        self.control_frame.pack(side=tk.RIGHT, padx=10, pady=10)
        
        # 创建预测结果标签
        self.result_label = tk.Label(self.control_frame, text="预测结果: ", font=("Arial", 24))
        self.result_label.pack(pady=10)
        
        self.confidence_label = tk.Label(self.control_frame, text="置信度: ", font=("Arial", 16))
        self.confidence_label.pack(pady=5)
        
        # 创建按钮
        self.predict_btn = tk.Button(self.control_frame, text="预测", command=self.predict_digit, font=("Arial", 14))
        self.predict_btn.pack(pady=5)
        
        self.clear_btn = tk.Button(self.control_frame, text="清除", command=self.clear_canvas, font=("Arial", 14))
        self.clear_btn.pack(pady=5)
        
        self.save_btn = tk.Button(self.control_frame, text="保存", command=self.save_image, font=("Arial", 14))
        self.save_btn.pack(pady=5)
        
        # 创建画笔粗细控制
        self.thickness_label = tk.Label(self.control_frame, text="画笔粗细:", font=("Arial", 12))
        self.thickness_label.pack(pady=5)
        
        self.thickness_var = tk.IntVar(value=15)
        self.thickness_scale = tk.Scale(self.control_frame, from_=5, to=30, orient=tk.HORIZONTAL, 
                                       variable=self.thickness_var)
        self.thickness_scale.pack(pady=5)
        
        # 创建概率条形图
        self.fig, self.ax = plt.subplots(figsize=(4, 3))
        self.canvas_widget = self.embed_matplotlib(self.fig)
        self.canvas_widget.pack(pady=10)
        
        # 初始化概率图
        self.probabilities = np.zeros(10)
        self.update_probability_chart()
        
        # 上次绘制位置
        self.last_x = None
        self.last_y = None
        
        # 保存图像计数
        self.save_count = 0
    
    def embed_matplotlib(self, fig):
        """将Matplotlib图形嵌入Tkinter窗口"""
        from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
        canvas = FigureCanvasTkAgg(fig, master=self.control_frame)
        canvas.draw()
        return canvas.get_tk_widget()
    
    def paint(self, event):
        """处理鼠标拖动绘制"""
        x, y = event.x, event.y
        r = self.thickness_var.get()  # 获取画笔粗细
        
        # 在Canvas上绘制
        self.canvas.create_oval(x-r, y-r, x+r, y+r, fill="white", outline="white")
        
        # 在PIL图像上绘制
        self.draw.ellipse([x-r, y-r, x+r, y+r], fill=255)
        
        # 如果有上一个点，则连接两点
        if self.last_x is not None and self.last_y is not None:
            self.canvas.create_line(self.last_x, self.last_y, x, y, width=r*2, fill="white")
            self.draw.line([self.last_x, self.last_y, x, y], fill=255, width=r*2)
        
        self.last_x = x
        self.last_y = y
    
    def clear_canvas(self):
        """清除画布"""
        self.canvas.delete("all")
        self.image = Image.new("L", (self.canvas_size, self.canvas_size), color=0)
        self.draw = ImageDraw.Draw(self.image)
        self.last_x = None
        self.last_y = None
        self.result_label.config(text="预测结果: ")
        self.confidence_label.config(text="置信度: ")
        self.probabilities = np.zeros(10)
        self.update_probability_chart()
    
    def save_image(self):
        """保存当前绘制的图像"""
        if not os.path.exists("drawn_digits"):
            os.makedirs("drawn_digits")
        
        filename = f"drawn_digits/digit_{self.save_count}.png"
        self.image.save(filename)
        print(f"图像已保存为 {filename}")
        self.save_count += 1
    
    def predict_digit(self):
        """预测绘制的数字"""
        # 调整图像大小为28x28
        img = self.image.resize((28, 28), Image.LANCZOS)
        
        # 转换为PyTorch张量
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        img_tensor = transform(img).unsqueeze(0)
        
        # 预测
        img_tensor = img_tensor.to(self.device)
        with torch.no_grad():
            output = self.model(img_tensor)
            probs = F.softmax(output, dim=1)[0].cpu().numpy()
            pred = output.argmax(dim=1, keepdim=True).item()
        
        # 更新UI
        self.result_label.config(text=f"预测结果: {pred}")
        self.confidence_label.config(text=f"置信度: {probs[pred]:.4f}")
        
        # 更新概率图表
        self.probabilities = probs
        self.update_probability_chart()
    
    def update_probability_chart(self):
        """更新概率条形图"""
        self.ax.clear()
        bars = self.ax.bar(range(10), self.probabilities)
        
        # 突出显示最高概率
        if np.max(self.probabilities) > 0:
            max_idx = np.argmax(self.probabilities)
            bars[max_idx].set_color('red')
        
        self.ax.set_xticks(range(10))
        self.ax.set_ylim([0, 1])
        self.ax.set_title("数字预测概率")
        self.fig.tight_layout()
        self.fig.canvas.draw()

def load_model(model_path):
    """加载训练好的模型"""
    # 检查CUDA是否可用
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # 加载模型
    model = Net().to(device)
    
    try:
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()
        print(f"模型已从 {model_path} 加载")
    except Exception as e:
        print(f"加载模型时出错: {e}")
        print("使用未训练的模型")
    
    return model, device

def main():
    model_path = "mnist_cnn.pt"
    
    # 如果模型文件不存在，提示用户
    if not os.path.exists(model_path):
        print(f"警告: 找不到模型文件 '{model_path}'")
        print("将使用未训练的模型。建议先运行 train.py 训练模型。")
    
    # 加载模型
    model, device = load_model(model_path)
    
    # 创建Tkinter应用
    root = tk.Tk()
    app = DrawingApp(root, model, device)
    root.mainloop()

if __name__ == "__main__":
    main() 